#include "ForwardPlusProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Game/GameState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"
#include "Game/Component/MultiModel.h"
#include <glm/gtc/random.hpp>



void ForwardPlusProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mSceneEntity = static_cast<LocalEntity*>(mGameScene->GetEntityByName("scene"));

	mPushConstants.debugIndex = 0;

	MultiModel& model = mSceneEntity->GetComponent<MultiModel>();
	mObjectSize = model.mVertexGroups.size();



	mPointLights.resize(mLocalScene->mLightNum);
	for (int32 i = 0; i < mLocalScene->mLightNum; ++i)
	{
		glm::vec3 color;
		do {
			color = glm::linearRand(glm::vec3(0, 0, 0), glm::vec3(1, 1, 1));
		} while (color.length() < 0.8f);
		mPointLights[i].intensity = color;
		mPointLights[i].radius = mLocalScene->mLightRadius;
		mPointLights[i].pos = glm::linearRand(mLocalScene->mMinLightPos, mLocalScene->mMaxLightPos);
	}




	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{


			mPointLightsBuffer = CustomBuffer::Builder()
				.SetData(mPointLights.data())
				.SetSize(sizeof(PointLight) * mLocalScene->mLightNum)
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "pointLights");




			mTileCountPerRow = (GRenderState->GetWidth() - 1) / TILE_SIZE + 1;
			mTileCountPerCol = (GRenderState->GetHeight() - 1) / TILE_SIZE + 1;

			mPushConstants.viewportSize.x = GRenderState->GetWidth();
			mPushConstants.viewportSize.y = GRenderState->GetHeight();
			mPushConstants.tileNum.x = mTileCountPerRow;
			mPushConstants.tileNum.y = mTileCountPerCol;


			mTileLightVisiblitiesBuffer = CustomBuffer::Builder()
				.SetUsage(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
				.SetSize(sizeof(LightVisiblity) * mTileCountPerRow * mTileCountPerCol)
				.Build(this, "tileLightVisiblities");



			RefCountPtr<RHIBuffer> cameraBuffer = GetBuffer("PreDepthProcess", "cameraUbo");

			RefCountPtr<RHITexture> emptyImage = GetTexture("PreDepthProcess", "emptyImage");

			RefCountPtr<RHISampler> emptySampler = GetSampler("PreDepthProcess", "emptySampler");

			RefCountPtr<RHITexture> depthImage = GetTexture("PreDepthProcess", "depthImage");

			RefCountPtr<RHISampler> depthSampler = RHISampler::Builder().Build();


			mDescriptor = RHIDescriptor::Builder()
				.PushShader<ForwardPlusVertexShader>()
				.PushShader<ForwardPlusFragmentShader>()
				.SetDescriptorCount(0, mObjectSize)
				.Build();

			for (int32 i = 0; i < mObjectSize; ++i)
			{
				mDescriptor->SetValue("cameraUbo", cameraBuffer);
				mDescriptor->SetValue("objectUbo", mSceneEntity->mObjectBuffer);
				mDescriptor->SetValue("materialUbo", mSceneEntity->mMaterialBuffer);
				mDescriptor->SetValue("tileLightVisiblities", mTileLightVisiblitiesBuffer);
				mDescriptor->SetValue("pointLights", mPointLightsBuffer);


				mDescriptor->SetValue("depthImage", depthImage, depthSampler);


				if (mSceneEntity->mMaterial.hasAlbedoMap == 1)
					mDescriptor->SetValue("albedoImage", mSceneEntity->mAlbedoImage, mSceneEntity->mAlbedoSampler);
				else
					mDescriptor->SetValue("albedoImage", emptyImage, emptySampler);


				if (mSceneEntity->mMaterial.hasNormalMap == 1)
					mDescriptor->SetValue("normalImage", mSceneEntity->mNormalImage, mSceneEntity->mNormalSampler);
				else
					mDescriptor->SetValue("normalImage", emptyImage, emptySampler);


				mSceneEntity->mDescriptorMasks[i] = mDescriptor->AllocateDescriptor();
			}

			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "colorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.Build(this, "depthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.SetPolygonMode(VK_POLYGON_MODE_FILL)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void ForwardPlusProcess::OnTick(float deltaTime)
{
	for (int i = 0; i < mLocalScene->mLightNum; i++) {
		mPointLights[i].pos += glm::vec3(0, 3.f, 0) * GGameState->mGameTimer.GetDeltaTime();
		if (mPointLights[i].pos.y > mLocalScene->mMaxLightPos.y) {
			mPointLights[i].pos.y -= (mLocalScene->mMaxLightPos.y - mLocalScene->mMinLightPos.y);
		}
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this](VkCommandBuffer commandBuffer)
		{
			mPointLightsBuffer->WriteMemory(mPointLights.data());


			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (int32 i = 0; i < mObjectSize; ++i)
			{

				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mSceneEntity->mDescriptorMasks[i].GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = mSceneEntity->mVertexBuffers[i]->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mSceneEntity->mIndexBuffers[i]->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &mPushConstants);

				vkCmdDrawIndexed(commandBuffer, mSceneEntity->mIndexBuffers[i]->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void ForwardPlusProcess::OnResize()
{
	mTileCountPerRow = (GRenderState->GetWidth() - 1) / TILE_SIZE + 1;
	mTileCountPerCol = (GRenderState->GetHeight() - 1) / TILE_SIZE + 1;

	mPushConstants.viewportSize.x = GRenderState->GetWidth();
	mPushConstants.viewportSize.y = GRenderState->GetHeight();
	mPushConstants.tileNum.x = mTileCountPerRow;
	mPushConstants.tileNum.y = mTileCountPerCol;

	mTileLightVisiblitiesBuffer = CustomBuffer::Builder()
		.SetUsage(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
		.SetSize(sizeof(LightVisiblity) * mTileCountPerRow * mTileCountPerCol)
		.Build(this, "tileLightVisiblities");

	RefCountPtr<RHITexture> depthImage = GetTexture("PreDepthProcess", "depthImage");

	RefCountPtr<RHISampler> depthSampler = RHISampler::Builder().Build();

	for (int32 i = 0; i < mObjectSize; ++i)
	{
		mDescriptor->SetValue("tileLightVisiblities", mTileLightVisiblitiesBuffer);

		mDescriptor->SetValue("depthImage", depthImage, depthSampler);

		mDescriptor->UpdateDescriptor(mSceneEntity->mDescriptorMasks[i]);
	}


	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "colorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.Build(this, "depthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
